import numpy as np
from RL.多臂老虎机.Solver import Solver


"""
__init__
    estimates: 期望每次使用老虎机每次都是中奖, 数组里面全部是1
    total_count: 初始化总次数

run_one_step
    运行一次总次数+1
    如果随机数小于 1/总次数, 随机选择一个老虎机, 也就是epsilon开始会很大, 按照总次数的增加会进行衰减
    否则, k=选择期望值最大的老虎机, 我们初始化的时候期望值全部是1, 所以第一次返回的一定是1
    r(奖励) = BernoulliBandit.step(K), 当前选择的老虎机的奖励值
    更新本次老虎机的estimates期望值估值
    返回当前操作的老虎机
    
"""
class DecayingEpsilonGreedy(Solver):
    """ epsilon值随时间衰减的epsilon-贪婪算法,继承Solver类 """
    def __init__(self, bandit, init_prob=1.0):
        super(DecayingEpsilonGreedy, self).__init__(bandit)
        self.estimates = np.array([init_prob] * self.bandit.K)
        self.total_count = 0

    def run_one_step(self):
        self.total_count += 1
        if np.random.random() < 1 / self.total_count:  # epsilon值随时间衰减
            k = np.random.randint(0, self.bandit.K)
        else:
            k = np.argmax(self.estimates)

        r = self.bandit.step(k)
        self.estimates[k] += 1. / (self.counts[k] + 1) * (r - self.estimates[k])

        return k
